from functools import wraps
from typing import Callable

from fastapi import Request
from app import constants
from app.core.log import logger
from app.core.redis import redis_client
from app.enums.limit import LimitTypeEnum
from app.exception import LimitException
from app.utils import context


def limit(
    name: str = "",
    key: str = "",
    prefix: str = "limit",
    period: int = 60,
    count: int = 10,
    limit_type: LimitTypeEnum = LimitTypeEnum.IP,
    describe: str = constants.LIMIT_DESCRIBE,
    enabled: bool = True,
):
    """
    限流装饰器

    Example:
        >>> @router.get("/captcha", summary="获取验证码")
        >>> @limit(name="验证码限流 （同一个ip，1分钟5次请求）", period=60, count=5, describe="您的操作过于频繁，请稍后再试")
        >>> async def captcha(request: Request) -> ApiResult[CaptchaVO]
        >>>     pass          # ↑ 使用@limit注解，需要有request参数

    Note:
        - 使用limit注解，函数必须是async协程函数
        - 使用limit注解，必须在方法的参数中添加`request: Request`

    :param name: 限流名称，用作备注
    :param key: 限流key，为空时自动使用方法名
    :param prefix: key前缀，最终的redis限流key = prefix + ":" + 处理后的key
    :param period: 给定的时间范围(秒)
    :param count: 一定时间内最多访问次数
    :param limit_type: 限流类型(ip或用户id)
    :param describe: 限流后返回的描述信息
    :param enabled: 是否启用限流
    :return: 装饰器函数
    """

    def decorator(func: Callable) -> Callable:
        """
        装饰器函数

        :param func: 被装饰的函数
        :return: 包装后的函数
        """

        @wraps(func)
        async def wrapper(*args, **kwargs):
            """
            包装函数，实现限流逻辑

            :param args: 位置参数
            :param kwargs: 关键字参数
            :return: 被装饰函数的返回值
            """
            # 如果未启用限流，直接执行原方法
            if not enabled:
                return await func(*args, **kwargs)

            # 获取fastApi request对象
            request = _get_request(*args, **kwargs)
            if not request:
                logger.warning(f"无法获取request对象，跳过对 {func.__name__} 的限流检查")
                return await func(*args, **kwargs)

            # 获取方法名
            method_name = key if key else func.__name__

            # 根据限流类型生成redis key
            redis_key = _get_redis_key(request, limit_type, prefix, method_name)

            # 执行限流检查
            if not await _check_limit(redis_key, period, count):
                raise LimitException(message=describe)

            # 执行原方法
            return await func(*args, **kwargs)

        return wrapper

    return decorator


def _get_request(*args, **kwargs) -> Request | None:
    """
    获取request对象

    :param args: 位置参数
    :param kwargs: 关键字参数
    :return: request对象
    """
    # 从args中获取request对象
    for arg in args:
        if isinstance(arg, Request):
            return arg

    # 从kwargs中获取request对象
    for key, value in kwargs.items():
        if isinstance(value, Request):
            return value
    return None


def _get_redis_key(request: Request, limit_type: LimitTypeEnum, prefix: str, method_name: str) -> str:
    """
    根据限流类型生成redis key

    :param request: 请求对象
    :param limit_type: 限流类型
    :param prefix: key前缀
    :param method_name: 方法名
    :return:
    """
    if limit_type == LimitTypeEnum.IP:
        # 获取客户端IP
        client_ip = _get_client_ip(request)
        return f"{prefix}:{method_name}:{client_ip}"
    elif limit_type == LimitTypeEnum.USER_ID:
        # 获取用户ID
        user_id = context.get_user_id()
        if not user_id or user_id == 0:
            raise LimitException(message="获取用户id失败")
        return f"{prefix}:{method_name}:{user_id}"
    else:
        raise LimitException(message="限流类型错误")


def _get_client_ip(request: Request) -> str:
    """
    获取客户端IP地址

    :param request: 请求对象
    :return: 客户端IP地址
    """
    # 尝试从X-Forwarded-For获取IP
    x_forwarded_for = request.headers.get("X-Forwarded-For")
    if x_forwarded_for:
        return x_forwarded_for.split(",")[0].strip()

    # 尝试从X-Real-IP获取IP
    x_real_ip = request.headers.get("X-Real-IP")
    if x_real_ip:
        return x_real_ip

    # 获取直接连接的客户端IP
    return request.client.host if request.client else "unknown"


async def _check_limit(redis_key: str, period: int, count: int) -> bool:
    """
    使用Lua脚本执行限流检查

    :param redis_key: Redis key
    :param period: 时间窗口(秒)
    :param count: 允许的最大请求数
    :return: True表示未触发限流，False表示触发限流
    """
    # Lua限流脚本
    lua_script = """
    local key_local = redis.call('setnx', KEYS[1], 0)
    if tonumber(key_local) == 0 then
        if tonumber(redis.call('get', KEYS[1])) >= tonumber(ARGV[2]) then
            return false
        else
            redis.call('incr', KEYS[1])
            return true
        end
    else
        redis.call('incr', KEYS[1])
        redis.call('expire', KEYS[1], ARGV[1])
        return true
    end
    """

    try:
        # 执行Lua脚本
        result = await redis_client.eval(
            lua_script,
            1,  # KEYS的数量
            redis_key,  # KEYS[1]
            str(period),  # ARGV[1]
            str(count)  # ARGV[2]
        )

        return bool(result)
    except Exception as e:
        # Redis执行异常时，默认触发限流
        logger.exception("限流检查异常", e)
        return False
